;  lzss compression library
;  copyright (c) 2004 Mateusz Tymek
;
;  17.03.2004
;  last updated 27.10.2004
;  version 1.01


uglobal

lzss:
  .rnode rd LZSS_DICTIONARY_SIZE + 1 + 0x100
  .lnode rd LZSS_DICTIONARY_SIZE + 1
  .parent rd LZSS_DICTIONARY_SIZE + 1
  sizeof.trees = $-lzss.rnode

  .dictionary rb LZSS_DICTIONARY_SIZE + LZSS_MAX_COUNT
  sizeof.dictionary = $-lzss.dictionary

  .best_match_len dd ?
  .best_match_pos dd ?
endg

iglobal
  lzss.packer_func dd 0
endg

; creates new empty LZSSData structure
proc lzssNewData
        invoke  GetProcessHeap
        invoke  HeapAlloc, eax,HEAP_ZERO_MEMORY,sizeof.LZSSData         ; alloc mem for LZSSData struc
        ret
endp


; destroys givem LZSSData structure
proc lzssFreeData .lzss_data
        push    eax ebx
        invoke  GetProcessHeap
        mov     ebx,eax
        mov     eax,[.lzss_data]
        invoke  HeapFree, ebx,0,[eax+LZSSData.data]
        invoke  HeapFree, ebx,0,[.lzss_data]
        pop     ebx eax
        ret
endp

;-----------------------------------------------------------------------------------------------


proc lzss_new_tree
        push    eax ecx esi edi
        mov     eax,LZSS_NO_TREE
        mov     ecx,LZSS_DICTIONARY_SIZE + 1
        mov     esi,lzss.lnode
        mov     edi,lzss.parent
     .clear:
        mov     [esi],eax
        mov     [edi],eax
        add     edi,4
        add     esi,4
        dec     ecx
        jnz     .clear
        mov     edi,lzss.rnode
        mov     ecx,LZSS_DICTIONARY_SIZE + 0x101
     .clear_btrees:
        mov     [edi],eax
        add     edi,4
        dec     ecx
        jnz     .clear_btrees
        pop     edi esi ecx eax
        ret
endp


proc lzss_insert_node .buffer_pos    ; inserts string at text_buf+[buf_pos] into the tree
        push    eax ebx ecx edx esi edi
        xor     eax,eax
        mov     [lzss.best_match_len],eax
        mov     [lzss.best_match_pos],eax

        mov     eax,[.buffer_pos]
        mov     edx,eax
        add     eax,lzss.dictionary
        movzx   eax,byte[eax]
        lea     ebx,[LZSS_DICTIONARY_SIZE+1+eax]
        mov     edx,1           ; edx will hold comparision result
        mov     eax,LZSS_NO_TREE
        mov     [lzss.rnode+edx*4],eax
        mov     [lzss.lnode+edx*4],eax

  .search_tree:
        or      edx,edx
        js      .left_node
     .right_node:
        cmp     [lzss.rnode+ebx*4],LZSS_NO_TREE
        je      .put_right

        mov     ebx,[lzss.rnode+ebx*4]
        jmp     .node_ok
     .put_right:
        mov     eax,[.buffer_pos]
        mov     [lzss.rnode+ebx*4],eax
        mov     [lzss.parent+eax*4],ebx
        jmp     .finish
     .left_node:
        cmp     [lzss.lnode+ebx*4],LZSS_NO_TREE
        je      .put_left

        mov     ebx,[lzss.lnode+ebx*4]
        jmp     .node_ok
     .put_left:
        mov     eax,[.buffer_pos]
        mov     [lzss.lnode+ebx*4],eax
        mov     [lzss.parent+eax*4],ebx
        jmp     .finish
     .node_ok:
        push    esi edi
        lea     edi,[lzss.dictionary+ebx]
        mov     esi,[.buffer_pos]
        lea     esi,[lzss.dictionary+esi]
        mov     ecx,LZSS_MAX_COUNT-1
        inc     esi     ;  we know that first character
        inc     edi     ;  already matches
        xor     edx,edx
        xor     eax,eax
     .compare:
        mov     dl,[esi]
        mov     al,byte[edi]
        sub     edx,eax
        or      edx,edx
        jnz     .compared
        inc     edi
        inc     esi
        dec     ecx
        jnz     .compare
     .compared:
        pop     esi edi

        mov     eax,LZSS_MAX_COUNT-1
        dec     ecx
        sub     eax,ecx
        cmp     eax,[lzss.best_match_len]
        jbe     .search_tree
        mov     [lzss.best_match_len],eax
        mov     [lzss.best_match_pos],ebx

        cmp     eax,LZSS_MAX_COUNT
        jae     .got_max_count

        jmp     .search_tree

     .got_max_count:
        mov     eax,[.buffer_pos]
        mov     edx,[lzss.parent+ebx*4]
        mov     [lzss.parent+eax*4],edx

        mov     edx,[lzss.lnode+ebx*4]
        mov     [lzss.lnode+eax*4],edx
        mov     [lzss.parent+edx*4],eax

        mov     edx,[lzss.rnode+ebx*4]
        mov     [lzss.rnode+eax*4],edx
        mov     [lzss.parent+edx*4],eax

        mov     edx,[lzss.parent+ebx*4]
        cmp     [lzss.rnode+edx*4],ebx
        je      .right
     .left:
        mov     [lzss.lnode+edx*4],eax
        jmp     .ok
     .right:
        mov     [lzss.rnode+edx*4],eax
     .ok:
        mov     eax,LZSS_NO_TREE
        mov     [lzss.parent+ebx*4],eax
        mov     [lzss.lnode+ebx*4],eax
        mov     [lzss.rnode+ebx*4],eax

     .finish:
        pop     edi esi edx ecx ebx eax
        ret
endp


; removes tree node at [buffer_pos]
proc lzss_delete_node .buffer_pos
        push    eax ebx ecx edx
        mov     eax,[.buffer_pos]
        cmp     dword[lzss.parent+eax*4],LZSS_NO_TREE   ; empty?
        je      .finish

        cmp     dword[lzss.rnode+eax*4],LZSS_NO_TREE
        je      .empty_right
        cmp     dword[lzss.lnode+eax*4],LZSS_NO_TREE
        je      .empty_left


        ; if we are here then none of children is empty

        mov     ebx,[lzss.lnode+eax*4]
        cmp     [lzss.rnode+ebx*4],LZSS_NO_TREE
        je      .got_free_rnode
      .find_free_rnode:
        mov     ebx,[lzss.rnode+ebx*4]
        cmp     [lzss.rnode+ebx*4],LZSS_NO_TREE
        jne     .find_free_rnode

        mov     edx,[lzss.parent+ebx*4]
        mov     ecx,[lzss.lnode+ebx*4]
        mov     [lzss.rnode+edx*4],ecx
        mov     [lzss.parent+ecx*4],edx

        mov     edx,[lzss.lnode+eax*4]
        mov     [lzss.lnode+ebx*4],edx
        mov     [lzss.parent+edx*4],ebx

  .got_free_rnode:

        mov     edx,[lzss.rnode+eax*4]
        mov     [lzss.rnode+ebx*4],edx
        mov     [lzss.parent+edx*4],ebx

        jmp     .remove

  .empty_left:
        mov     ebx,[lzss.rnode+eax*4]
        jmp     .remove
  .empty_right:
        mov     ebx,[lzss.lnode+eax*4]

  .remove:

        mov     edx,[lzss.parent+eax*4]         ; now edx holds parent of node being removed
        mov     [lzss.parent+ebx*4],edx

        ; now fix the valid child
        cmp     [lzss.rnode+edx*4],eax
        je      .fix_right_child
  .fix_left_child:
        mov     [lzss.lnode+edx*4],ebx
        jmp     .clear
  .fix_right_child:
        mov     [lzss.rnode+edx*4],ebx

  .clear:       ; clear the node that is being removed
        mov     [lzss.parent+eax*4],LZSS_NO_TREE
        mov     [lzss.rnode+eax*4],LZSS_NO_TREE
        mov     [lzss.lnode+eax*4],LZSS_NO_TREE

  .finish:

        pop     edx ecx ebx eax
        ret
endp


proc lzssPackData .buffer,.buffer_size
local .lzss_data:DWORD
local .bitstream:DWORD

        push    ebx ecx edx esi edi
        ; fill the dictionary with zeroes
        mov     edi,lzss.dictionary
        mov     ecx,(LZSS_DICTIONARY_SIZE + LZSS_MAX_COUNT)/2
        xor     eax,eax
      .clear_dictionary:
        mov     [edi],ax
        add     edi,2
        dec     ecx
        jnz     .clear_dictionary

        stdcall lzssNewData
        or      eax,eax
        jz      .finish
        mov     [.lzss_data],eax
        mov     edi,eax
        invoke  GetProcessHeap
        mov     ebx,eax
        mov     eax,[.buffer_size]       ; calculate maximum memory that can be needed by
        shr     eax,3                   ; compressed data. It is data_size + data_size/8 ...
        inc     eax                     ; and one more byte just to be sure :)
        add     eax,[.buffer_size]       ;
        invoke  HeapAlloc, ebx,HEAP_ZERO_MEMORY,eax
        or      eax,eax
        jz      .finish
        mov     [edi+LZSSData.data],eax
        mov     ecx,[.buffer_size]
        mov     [edi+LZSSData.unpacked_size],ecx
        mov     eax,[.buffer]            ; esi will point to input buffer
        add     eax,ecx         ; add eax,[buffer_size]
        stdcall BitsCreateFrom, [edi+LZSSData.data]
        mov     [.bitstream],eax

        ; edi will point to the end of the dictionary...
        mov     edi,LZSS_DICTIONARY_SIZE-LZSS_MAX_COUNT
        ; ... and esi will mark the start
        xor     esi,esi
        stdcall lzss_new_tree


        xor     edx,edx
     .read_first_bytes:
        mov     eax,[.buffer]
        mov     al,[eax]
        inc     [.buffer]
        mov     [lzss.dictionary+edi+edx],al
        inc     edx
        cmp     edx,LZSS_MAX_COUNT
        jae     .got_first_bytes
        dec     [.buffer_size]
        jz      .got_first_bytes
        jmp     .read_first_bytes
     .got_first_bytes:

        or      edx,edx
        jz      .packed

        mov     ecx,LZSS_MAX_COUNT-1
        mov     eax,edi
        dec     eax
     .prepare_tree:
        stdcall lzss_insert_node,eax
        dec     eax
        dec     ecx
        jnz     .prepare_tree
        stdcall lzss_insert_node,edi         ; insert the string just read; best_match_len and best_match_pos are set

     .packer_loop:
        cmp     [lzss.best_match_len],edx    ; best_match_len may be too long near the end
        jbe      .lzss.best_match_len_ok
        mov     [lzss.best_match_len],edx
     .lzss.best_match_len_ok:

        cmp     [lzss.best_match_len],LZSS_THRESHOLD   ; send unpacked byte or ofs-length pair?
        ja      .send_packed
        stdcall BitsPut, [.bitstream],0,1       ; send "unpacked data" flag
        mov     al,[lzss.dictionary+edi]
        stdcall BitsPut, [.bitstream],eax,8     ; send byte directly
        mov     [lzss.best_match_len],1              ; mark that we have processed just one byte

        jmp     .item_packed
     .send_packed:
        stdcall BitsPut, [.bitstream],1,1       ; "packed data" flag
        stdcall BitsPut, [.bitstream],[lzss.best_match_pos],LZSS_OFFSET_BITS
        mov     eax,[lzss.best_match_len]
        sub     eax,LZSS_THRESHOLD+1
        stdcall BitsPut, [.bitstream],eax,LZSS_COUNT_BITS


     .item_packed:
        xor     ecx,ecx
        mov     ebx,[lzss.best_match_len]    ; ebx will hold value of "last match length"
        cmp     [.buffer_size],0
        je      .new_bytes_ok
     .get_new_bytes:
        cmp     ecx,ebx                 ; compare ecx to [last_match_len]
        jae     .new_bytes_ok
        dec     [.buffer_size]
        jz      .new_bytes_ok
        stdcall lzss_delete_node,esi
        mov     eax,[.buffer]
        inc     [.buffer]
        mov     al,[eax]
        mov     [lzss.dictionary+esi],al        ; put new byte in the dictionary
        cmp     esi,LZSS_MAX_COUNT
        jae     .no_extra
        mov     [lzss.dictionary+esi+LZSS_DICTIONARY_SIZE],al
     .no_extra:
        inc     esi
        inc     edi
        mov     eax,LZSS_DICTIONARY_SIZE
        dec     eax
        and     esi,eax
        and     edi,eax

        stdcall lzss_insert_node,edi         ; insert new string

        inc     ecx
        jmp     .get_new_bytes
     .new_bytes_ok:
        cmp     ecx,ebx
        jae     .packed
        inc     ecx
        stdcall lzss_delete_node,esi
        inc     esi
        inc     edi
        mov     eax,LZSS_DICTIONARY_SIZE
        dec     eax
        and     esi,eax
        and     edi,eax
        dec     edx
        jz      .new_bytes_ok

        stdcall lzss_insert_node,edi
        jmp     .new_bytes_ok
     .packed:

        or      edx,edx
        jnz     .packer_loop

        mov     ebx,[.lzss_data]
        stdcall BitsGetSize, [.bitstream]
        mov     [ebx+LZSSData.packed_size],eax
     .finish:

        mov     eax,[.lzss_data]
        pop     edi esi edx ecx ebx
        ret
endp


;iglobal
;  szOutFile db "out.dat", 0
;  szOutFile2 db "out2.dat", 0
;endg
;

;-----------------------------------------------------------------------------------------------



; debug lines
;iglobal
;  szDictionary db "ofs: %d, len: %d", 13, 10, 0
;  szDirect     db "direct byte: %d ('%c')", 13, 10, 0
;  szEndMarker  db "eof", 13, 10, 0
;endg


; unpacks given LZSSData structure and returns pointer
; when you have finished with this data, you should
; free this memory using HeapFree()
proc lzssUnpackData .lzss_data
local .uncompressed:DWORD
local .size:DWORD
        push    ebx ecx edx esi edi

        ; fill the dictionary with zeroes
        mov     edi,lzss.dictionary
        mov     ecx,(LZSS_DICTIONARY_SIZE + LZSS_MAX_COUNT)/2
        xor     eax,eax
      .clear_dictionary:
        mov     [edi],ax
        add     edi,2
        dec     ecx
        jnz     .clear_dictionary

        mov     ebx,[.lzss_data]
        invoke  GetProcessHeap
        invoke  HeapAlloc, eax,HEAP_ZERO_MEMORY,[ebx+LZSSData.unpacked_size]
        mov     [.uncompressed],eax
        mov     edi,eax
        stdcall BitsCreateFrom, [ebx+LZSSData.data]
        mov     esi,eax
        mov     eax,[ebx+LZSSData.unpacked_size]
        mov     [.size],eax
        mov     edx,LZSS_DICTIONARY_SIZE - LZSS_MAX_COUNT

   .unpack:
        stdcall BitsGet,esi,1
        or      eax,eax
        jnz     .from_dictionary
        stdcall BitsGet,esi,8

     ; debug lines
;     pusha
;     cinvoke printf, szDirect, eax,eax
;     popa
        mov     [edi],al
        inc     edi

        mov     [lzss.dictionary+edx],al
        inc     edx
        mov     eax,LZSS_DICTIONARY_SIZE
        dec     eax
        and     edx,eax
        dec     [.size]
        jz      .finish
        jmp     .unpack


    .from_dictionary:
        stdcall BitsGet, esi,LZSS_OFFSET_BITS         ; load position
;     push    eax  ; debug line
        mov     ebx,eax    ; no, unpack portion
        stdcall BitsGet, esi,LZSS_COUNT_BITS          ; load counter
        mov     ecx,eax
        add     ecx,LZSS_THRESHOLD+1
        sub     [.size],ecx

; debug lines
;     pop     eax
;     pusha
;     cinvoke printf, szDictionary, eax,ecx
;     popa

     .copy:
        mov     al,[lzss.dictionary+ebx]
        mov     [edi],al
        mov     [lzss.dictionary+edx],al
        inc     ebx
        inc     edi
        inc     edx
        mov     eax,LZSS_DICTIONARY_SIZE
        dec     eax
        and     edx,eax
        and     ebx,eax

        dec     ecx
        jnz     .copy
        cmp     [.size],0

        je      .finish
        jmp     .unpack
     .finish:

;  debug lines
;     pusha
;     cinvoke printf, szEndMarker
;     popa

        stdcall BitsDestroy, esi


;        invoke  CreateFile, szOutFile2,GENERIC_WRITE,0,0,CREATE_ALWAYS,0,0
;        mov     ebx,eax
;        invoke  WriteFile,eax,lzss.dictionary,sizeof.dictionary,tmp,0
;        invoke  CloseHandle,ebx
;

        mov     eax,[.uncompressed]


        pop     edi esi edx ecx ebx
        ret
endp



; file routines

; compress data and store it in given file; returns 0 on fail
proc lzssSaveFile .file_name,.buffer,.size
local .hdr:LZSSFileHeader
local .tmp:DWORD
        push    ebx ecx edx esi
        mov     eax,[lzss.packer_func]
        or      eax,eax
        jnz     .packer_func_ok
        mov     [lzss.packer_func],lzssPackData
     .packer_func_ok:
        stdcall [lzss.packer_func], [.buffer],[.size]
        mov     esi,eax
        mov     [.hdr.id],"LZSS"
        mov     eax,[.size]
        mov     [.hdr.unpacked_size],eax
        invoke  CreateFile, [.file_name],GENERIC_WRITE,0,0,CREATE_ALWAYS,0,0
        cmp     eax,INVALID_HANDLE_VALUE
        je      .error
        mov     ebx,eax
        lea     eax,[.hdr]
        lea     edx,[.tmp]
        invoke  WriteFile,ebx,eax,sizeof.LZSSFileHeader,edx,0
        mov     ecx,[esi+LZSSData.packed_size]
        mov     eax,[esi+LZSSData.data]
        invoke  WriteFile,ebx,eax,ecx,tmp,0
        invoke  CloseHandle,ebx
        stdcall lzssFreeData,esi
        xor     eax,eax
        dec     eax
     .finish:
        pop     esi edx ecx ebx
        ret
     .error:
        xor     eax,eax
        jmp     .finish
endp


; loads packed data from given file, unpacks it and returns pointer
; (0 on fail); ptr_size should be a pointer to dword variable
; that will contain unpacked data size
proc lzssLoadFile .file_name, .ptr_size
local .input_buf:DWORD
local .input_size:DWORD
        push    ebx ecx edx esi edi
        invoke  CreateFile, [.file_name],GENERIC_READ,0,0,OPEN_EXISTING,0,0
        cmp     eax,INVALID_HANDLE_VALUE
        je      .error
        mov     ebx,eax
        invoke  SetFilePointer,ebx,0,0,FILE_END
        mov     [.input_size],eax
        mov     esi,eax
        invoke  GetProcessHeap
        invoke  HeapAlloc, eax,HEAP_ZERO_MEMORY,esi
        or      eax,eax
        jz      .finish
        mov     [.input_buf],eax
        invoke  SetFilePointer, ebx,0,0,FILE_BEGIN
        invoke  ReadFile, ebx,[.input_buf],esi,tmp,0
        invoke  CloseHandle,ebx
        stdcall lzssNewData
        or      eax,eax
        jz      .finish
        mov     esi,eax
        mov     eax,[.input_buf]
        add     eax,sizeof.LZSSFileHeader
        mov     [esi+LZSSData.data],eax
        mov     eax,[.input_buf]
        mov     eax,[eax+LZSSFileHeader.unpacked_size]
        mov     [esi+LZSSData.unpacked_size],eax
        stdcall lzssUnpackData, esi
        or      eax,eax
        jz      .finish
        mov     ebx, eax

        mov     eax,[.ptr_size]
        or      eax,eax
        jz      .no_size_ptr
        mov     edx,[esi+LZSSData.unpacked_size]
        mov     [eax],edx
     .no_size_ptr:
        stdcall lzssFreeData,esi
        mov     eax,ebx
     .finish:
        pop     edi esi edx ecx ebx
        ret
     .error:
        xor     eax,eax
        jmp     .finish
endp